#include <CL/sycl.hpp>
#include <iostream>
#include <vector>

namespace sycl = cl::sycl;

class GaussElimination {
public:
  GaussElimination(const std::vector<float> &matrix, const std::vector<float> &rhs)
      : matrix(matrix), rhs(rhs) {}

  void solve() {
    sycl::queue queue(sycl::default_selector{});
    const int size = matrix.size();
    std::vector<float> result(size, 0.0f);

    {
      sycl::buffer<float> matrixBuf(matrix);
      sycl::buffer<float> rhsBuf(rhs);
      sycl::buffer<float> resultBuf(result);

      queue.submit([&](sycl::handler &cgh) {
        auto matrixAcc = matrixBuf.get_access<sycl::access::mode::read_write>(cgh);
        auto rhsAcc = rhsBuf.get_access<sycl::access::mode::read_write>(cgh);
        auto resultAcc =
            resultBuf.get_access<sycl::access::mode::discard_write>(cgh);

        cgh.parallel_for<class GaussEliminationKernel>(
            sycl::range<1>(size), [=](sycl::id<1> idx) {
              const int row = idx[0];

              for (int i = row + 1; i < size; ++i) {
                const float ratio = matrixAcc[i * size + row] /
                                    matrixAcc[row * size + row];

                for (int j = row; j < size; ++j) {
                  matrixAcc[i * size + j] -= ratio * matrixAcc[row * size + j];
                }

                rhsAcc[i] -= ratio * rhsAcc[row];
              }
            });
      });
    }

    queue.wait();

    std::cout << "Solution:" << std::endl;
    for (int i = 0; i < size; ++i) {
      std::cout << "x[" << i << "] = " << result[i] << std::endl;
    }
  }

private:
  const std::vector<float> &matrix;
  const std::vector<float> &rhs;
};

int main() {
  const int size = 4;
  const std::vector<float> matrix = {2.0f, 1.0f, -1.0f, 8.0f, 4.0f, -3.0f, -1.0f, 1.0f,
                                     -2.0f, 1.0f, 2.0f, 3.0f, 3.0f, -1.0f, 3.0f, -1.0f};
  const std::vector<float> rhs = {11.0f, -1.0f, 10.0f, -1.0f};

  GaussElimination solver(matrix, rhs);
  solver.solve();

  return 0;
}